//
// Created by 朱林 on 2020/9/21.
//
#include "multi_pose.h"

inline float clip(float value, float min, float max) {
    if (value < 0) {
        return 0;
    } else if (value > max) {
        return max;
    } else {
        return value;
    }
}


static void generatePointList(std::vector<float> &poseScores, std::vector<Point> pointList,
                              std::vector<std::vector<float>> &poseKeyPointScores,
                              std::vector<std::vector<CV::Point>> &poseKeyPointCoords) {
    const int poseCount = poseScores.size();
    for (int i = 0; i < poseCount; ++i) {
        if (poseScores[i] > MIN_POSE_SCORE) {
            for (int id = 0; id < NUM_KEYPOINTS; ++id) {
                if (poseKeyPointScores[i][id] > SCORE_THRESHOLD) {
                    CV::Point point = poseKeyPointCoords[i][id];
                    Point point1;
                    point1.x = point.fX;
                    point1.y = point.fY;
                    pointList.push_back(point1);
                }
            }
        }
    }
}


static CV::Point getCoordsFromTensor(const Tensor *dataTensor,
                                     int id, int x, int y,
                                     bool getCoord = true) {
    // dataTensor must be [1,c,h,w]
    auto dataPtr = dataTensor->host<float>();
    const int xOffset = dataTensor->channel() / 2;
    const int indexPlane = y * dataTensor->stride(2) + x;
    const int indexY = id * dataTensor->stride(1) + indexPlane;
    const int indexX = (id + xOffset) * dataTensor->stride(1) + indexPlane;
    CV::Point point;
    if (getCoord) {
        point.set(dataPtr[indexX], dataPtr[indexY]);
    } else {
        point.set(0.0, dataPtr[indexY]);
    }
    return point;
};

// decode pose and posenet model reference from https://github.com/rwightman/posenet-python
static int decodePoseImpl(float curScore, int curId, const CV::Point &originalOnImageCoords,
                          const Tensor *heatmaps,
                          const Tensor *offsets, const Tensor *displacementFwd,
                          const Tensor *displacementBwd,
                          std::vector<float> &instanceKeypointScores,
                          std::vector<CV::Point> &instanceKeypointCoords) {
    instanceKeypointScores[curId] = curScore;
    instanceKeypointCoords[curId] = originalOnImageCoords;
    const int height = heatmaps->height();
    const int width = heatmaps->width();
    std::map<std::string, int> poseNamesID;
    for (int i = 0; i < PoseNames.size(); ++i) {
        poseNamesID[PoseNames[i]] = i;
    }

    auto traverseToTargetKeypoint = [=](int edgeId, const CV::Point &sourcekeypointCoord,
                                        int targetKeypointId,
                                        const Tensor *displacement) {
        int sourceKeypointIndicesX =
                static_cast<int>(clip(round(sourcekeypointCoord.fX / (float) OUTPUT_STRIDE), 0,
                                      (float) (width - 1)));
        int sourceKeypointIndicesY =
                static_cast<int>(clip(round(sourcekeypointCoord.fY / (float) OUTPUT_STRIDE), 0,
                                      (float) (height - 1)));

        auto displacementCoord =
                getCoordsFromTensor(displacement, edgeId, sourceKeypointIndicesX,
                                    sourceKeypointIndicesY);
        float displacedPointX = sourcekeypointCoord.fX + displacementCoord.fX;
        float displacedPointY = sourcekeypointCoord.fY + displacementCoord.fY;

        int displacedPointIndicesX =
                static_cast<int>(clip(round(displacedPointX / OUTPUT_STRIDE), 0,
                                      (float) (width - 1)));
        int displacedPointIndicesY =
                static_cast<int>(clip(round(displacedPointY / OUTPUT_STRIDE), 0,
                                      (float) (height - 1)));

        float score =
                getCoordsFromTensor(heatmaps, targetKeypointId, displacedPointIndicesX,
                                    displacedPointIndicesY, false).fY;
        auto offset = getCoordsFromTensor(offsets, targetKeypointId, displacedPointIndicesX,
                                          displacedPointIndicesY);

        CV::Point imageCoord;
        imageCoord.fX = displacedPointIndicesX * OUTPUT_STRIDE + offset.fX;
        imageCoord.fY = displacedPointIndicesY * OUTPUT_STRIDE + offset.fY;

        return std::make_pair(score, imageCoord);
    };

    MNN_ASSERT((NUM_KEYPOINTS - 1) == PoseChain.size());

    for (int edge = PoseChain.size() - 1; edge >= 0; --edge) {
        const int targetKeypointID = poseNamesID[PoseChain[edge].first];
        const int sourceKeypointID = poseNamesID[PoseChain[edge].second];
        if (instanceKeypointScores[sourceKeypointID] > 0.0 &&
            instanceKeypointScores[targetKeypointID] == 0.0) {
            auto curInstance = traverseToTargetKeypoint(edge,
                                                        instanceKeypointCoords[sourceKeypointID],
                                                        targetKeypointID, displacementBwd);
            instanceKeypointScores[targetKeypointID] = curInstance.first;
            instanceKeypointCoords[targetKeypointID] = curInstance.second;
        }
    }

    for (int edge = 0; edge < PoseChain.size(); ++edge) {
        const int sourceKeypointID = poseNamesID[PoseChain[edge].first];
        const int targetKeypointID = poseNamesID[PoseChain[edge].second];
        if (instanceKeypointScores[sourceKeypointID] > 0.0 &&
            instanceKeypointScores[targetKeypointID] == 0.0) {
            auto curInstance = traverseToTargetKeypoint(edge,
                                                        instanceKeypointCoords[sourceKeypointID],
                                                        targetKeypointID, displacementFwd);
            instanceKeypointScores[targetKeypointID] = curInstance.first;
            instanceKeypointCoords[targetKeypointID] = curInstance.second;
        }
    }

    return 0;
}

static int decodeMultiPose(const Tensor *offsets, const Tensor *displacementFwd,
                           const Tensor *displacementBwd, const Tensor *heatmaps,
                           std::vector<float> &poseScores,
                           std::vector<std::vector<float>> &poseKeypointScores,
                           std::vector<std::vector<CV::Point>> &poseKeypointCoords,
                           CV::Point &scale) {
    // keypoint_id, score, coord((x,y))
    typedef std::pair<int, std::pair<float, CV::Point>> partsType;
    std::vector<partsType> parts;

    const int channel = heatmaps->channel();
    const int height = heatmaps->height();
    const int width = heatmaps->width();
    auto maximumFilter = [&parts, width, height](const int id, const float *startPtr) {
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                // check whether (y,x) is the max value around the neighborhood
                bool isMaxVaule = true;
                float maxValue = startPtr[y * width + x];
                {
                    for (int i = -LOCAL_MAXIMUM_RADIUS; i < (LOCAL_MAXIMUM_RADIUS + 1); ++i) {
                        for (int j = -LOCAL_MAXIMUM_RADIUS; j < (LOCAL_MAXIMUM_RADIUS + 1); ++j) {
                            float value = 0.0f;
                            int yCoord = y + i;
                            int xCoord = x + j;
                            if (yCoord >= 0 && yCoord < height && xCoord >= 0 && xCoord < width) {
                                value = startPtr[yCoord * width + xCoord];
                            }
                            if (maxValue < value) {
                                isMaxVaule = false;
                                break;
                            }
                        }
                    }
                }

                if (isMaxVaule && maxValue >= SCORE_THRESHOLD) {
                    CV::Point coord;
                    coord.set(x, y);
                    parts.push_back(std::make_pair(id, std::make_pair(maxValue, coord)));
                }
            }
        }
    };
    auto scoresPtr = heatmaps->host<float>();
    for (int id = 0; id < channel; ++id) {
        auto idScoresPtr = scoresPtr + id * width * height;
        maximumFilter(id, idScoresPtr);
    }
    // sort the parts according to score
    std::sort(parts.begin(), parts.end(),
              [](const partsType &a, const partsType &b) {
                  return a.second.first > b.second.first;
              });
    const int squareNMSRadius = NMS_RADIUS * NMS_RADIUS;
    auto withinNMSRadius = [=, &poseKeypointCoords](const CV::Point &point, const int id) {
        bool withinThisPointRadius = false;
        for (int i = 0; i < poseKeypointCoords.size(); ++i) {
            const auto &curPoint = poseKeypointCoords[i][id];
            const auto sum = powf((curPoint.fX - point.fX), 2) + powf((curPoint.fY - point.fY), 2);
            if (sum <= squareNMSRadius) {
                withinThisPointRadius = true;
                break;
            }
        }
        return withinThisPointRadius;
    };
    std::vector<float> instanceKeypointScores(NUM_KEYPOINTS);
    std::vector<CV::Point> instanceKeypointCoords(NUM_KEYPOINTS);
    auto getInstanceScore = [&]() {
        float notOverlappedScores = 0.0f;
        const int poseNums = poseKeypointCoords.size();
        if (poseNums == 0) {
            for (int i = 0; i < NUM_KEYPOINTS; ++i) {
                notOverlappedScores += instanceKeypointScores[i];
            }
        } else {
            for (int id = 0; id < NUM_KEYPOINTS; ++id) {
                if (!withinNMSRadius(instanceKeypointCoords[id], id)) {
                    notOverlappedScores += instanceKeypointScores[id];
                }
            }
        }
        return notOverlappedScores / NUM_KEYPOINTS;
    };
    int poseCount = 0;
    for (const auto &part : parts) {
        if (poseCount >= MAX_POSE_DETECTIONS) {
            break;
        }
        const auto curScore = part.second.first;
        const auto curId = part.first;
        const auto &curPoint = part.second.second;
        const auto offsetXY = getCoordsFromTensor(offsets, curId, (int) curPoint.fX,
                                                  (int) curPoint.fY);
        CV::Point originalOnImageCoords;
        originalOnImageCoords.fX = curPoint.fX * OUTPUT_STRIDE + offsetXY.fX;
        originalOnImageCoords.fY = curPoint.fY * OUTPUT_STRIDE + offsetXY.fY;

        if (withinNMSRadius(originalOnImageCoords, curId)) {
            continue;
        }
        ::memset(instanceKeypointScores.data(), 0, sizeof(float) * NUM_KEYPOINTS);
        ::memset(instanceKeypointCoords.data(), 0, sizeof(CV::Point) * NUM_KEYPOINTS);
        decodePoseImpl(curScore, curId, originalOnImageCoords, heatmaps, offsets, displacementFwd,
                       displacementBwd,
                       instanceKeypointScores, instanceKeypointCoords);

        float poseScore = getInstanceScore();
        if (poseScore > MIN_POSE_SCORE) {
            poseScores.push_back(poseScore);
            poseKeypointScores.push_back(instanceKeypointScores);
            poseKeypointCoords.push_back(instanceKeypointCoords);
            poseCount++;
        }
    }
    // scale the pose keypoint coords
    for (int i = 0; i < poseCount; ++i) {
        for (int id = 0; id < NUM_KEYPOINTS; ++id) {
            poseKeypointCoords[i][id].fX *= scale.fX;
            poseKeypointCoords[i][id].fY *= scale.fY;
        }
    }
    return 0;
}

multiPose::multiPose() {

}

multiPose::~multiPose() {
    pose_interpreter->releaseModel();
    pose_interpreter->releaseSession(pose_session);
}

bool multiPose::init(const std::string &mnn_path, int num_thread_, bool openCL) {
    num_thread = num_thread_;
    pose_interpreter = std::shared_ptr<MNN::Interpreter>(
            MNN::Interpreter::createFromFile(mnn_path.c_str()));
    MNN::ScheduleConfig netConfig;
    netConfig.type = MNN_FORWARD_CPU;
    netConfig.numThread = num_thread;
    pose_session = pose_interpreter->createSession(netConfig);
    input_tensor = pose_interpreter->getSessionInput(pose_session, nullptr);
    return true;
}

std::vector<Point> multiPose::detect(cv::Mat &img) {
    std::vector<Point> ret_list = {};
    if (img.empty()) {
        std::cout << "image is empty ,please check!" << std::endl;
        return ret_list;
    }
    int image_h = img.rows;
    int image_w = img.cols;
    const int target_w =
            static_cast<int>((float) image_w / (float) OUTPUT_STRIDE) * OUTPUT_STRIDE + 1;
    const int target_h =
            static_cast<int>((float) image_h / (float) OUTPUT_STRIDE) * OUTPUT_STRIDE + 1;
    CV::Point scale;
    scale.fX = (float) image_w / (float) target_w;
    scale.fY = (float) image_h / (float) target_h;

    if (input_tensor->elementSize() <= 4) {
        pose_interpreter->resizeTensor(input_tensor, {1, 3, target_h, target_w});
        pose_interpreter->resizeSession(pose_session);
    }
    {
        cv::Mat image;
        cv::resize(img, image, cv::Size(target_w, target_h));
        std::shared_ptr<MNN::CV::ImageProcess> pretreat(
                MNN::CV::ImageProcess::create(MNN::CV::BGR, MNN::CV::RGB, means, 3,
                                              norms, 3));
        pretreat->convert(image.data, target_w, target_h, image.step[0], input_tensor);
    }
    {
        AUTOTIME;
        pose_interpreter->runSession(pose_session);
    }
    auto offsets = pose_interpreter->getSessionOutput(pose_session, OFFSET_NODE_NAME);
    auto displacementFwd = pose_interpreter->getSessionOutput(pose_session, DISPLACE_FWD_NODE_NAME);
    auto displacementBwd = pose_interpreter->getSessionOutput(pose_session, DISPLACE_BWD_NODE_NAME);
    auto heat_maps = pose_interpreter->getSessionOutput(pose_session, HEATMAPS);

    Tensor offsetsHost(offsets, Tensor::CAFFE);
    Tensor displacementFwdHost(displacementFwd, Tensor::CAFFE);
    Tensor displacementBwdHost(displacementBwd, Tensor::CAFFE);
    Tensor heatMapsHost(heat_maps, Tensor::CAFFE);

    offsets->copyToHostTensor(&offsetsHost);
    displacementFwd->copyToHostTensor(&displacementFwdHost);
    displacementBwd->copyToHostTensor(&displacementBwdHost);
    heat_maps->copyToHostTensor(&heatMapsHost);
    std::vector<float> poseScores;
    std::vector<std::vector<float>> poseKeyPointScores;
    std::vector<std::vector<CV::Point>> poseKeyPointCoords;
    {
        AUTOTIME;
        decodeMultiPose(&offsetsHost, &displacementFwdHost, &displacementBwdHost, &heatMapsHost,
                        poseScores,
                        poseKeyPointScores, poseKeyPointCoords, scale);
    }
    generatePointList(poseScores, ret_list, poseKeyPointScores, poseKeyPointCoords);
}